import argparse
import logging
import os
import time
import signal
import sys
from typing import Dict, List, Optional

from config.settings import DEBUG
from core.analyzer import CodeAnalyzer
from core.graphdb import GraphDBManager
from core.generator import DocGenerator
from utils.ollama_manager import OllamaProcessManager

# 配置日志
logging.basicConfig(
    level=logging.DEBUG if DEBUG else logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 全局变量跟踪生成器和进程管理器
generator = None
ollama_manager = None

def signal_handler(sig, frame):
    """处理中断信号(Ctrl+C)"""
    print("\n捕获到中断信号，正在清理资源...")
    
    if generator and hasattr(generator, 'llm'):
        # 取消当前生成
        generator.llm.cancel_generation()
    
    if ollama_manager:
        # 终止所有活跃进程
        count = ollama_manager.terminate_all_processes()
        if count > 0:
            print(f"已终止 {count} 个活跃进程")
    
    print("已清理资源，退出程序")
    sys.exit(0)

def check_and_fix_environment():
    """检查并修复环境设置"""
    # 确保输出目录存在
    output_dir = os.path.join(os.getcwd(), "generated_docs")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        logger.info(f"创建输出目录: {output_dir}")
    
    # 初始化Ollama进程管理器
    global ollama_manager
    ollama_manager = OllamaProcessManager()
    
    # 确保Ollama服务正在运行
    if not ollama_manager.ensure_service_running():
        logger.warning("无法确认Ollama服务状态，尝试继续")

def main():
    """程序入口点"""
    # 设置信号处理
    signal.signal(signal.SIGINT, signal_handler)
    
    # 检查环境
    check_and_fix_environment()
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='代码文档生成系统')
    parser.add_argument('--repo', '-r', type=str, help='要分析的代码仓库路径', required=True)
    parser.add_argument('--file', '-f', type=str, help='要生成文档的特定文件路径')
    parser.add_argument('--function', '-func', type=str, help='要生成文档的特定函数名')
    parser.add_argument('--overview', '-o', action='store_true', help='生成项目概述文档')
    parser.add_argument('--all', '-a', action='store_true', help='为所有文件生成文档')
    parser.add_argument('--observe', '-obs', action='store_true', help='启用观察模式，显示推理过程')
    parser.add_argument('--subprocess', '-sp', action='store_true', help='使用子进程模式，提供更好的隔离')
    parser.add_argument('--check', '-c', action='store_true', help='仅检查环境和资源状态')
    parser.add_argument('--restart', '-rs', action='store_true', help='重启Ollama服务')
    
    args = parser.parse_args()
    
    # 特殊命令处理
    if args.restart:
        logger.info("正在重启Ollama服务...")
        success = ollama_manager.restart_service()
        if success:
            print("已成功重启Ollama服务")
        else:
            print("重启Ollama服务失败，请手动检查")
        return
    
    if args.check:
        # 打印系统资源状态
        print("系统资源状态:")
        resources = ollama_manager.check_resources()
        
        if "error" in resources:
            print(f"  检查失败: {resources['error']}")
        else:
            print(f"  内存: {resources['memory']['available']} 可用 / {resources['memory']['total']} 总计 ({resources['memory']['percent']} 使用)")
            
            if resources["gpu"]:
                print("  GPU状态:")
                for gpu in resources["gpu"]:
                    print(f"    GPU {gpu['id']}: {gpu['memory_used']} 使用 / {gpu['memory_total']} 总计 ({gpu['percent']} 使用)")
            else:
                print("  未检测到GPU或无法获取GPU信息")
        
        # 检查活跃进程
        processes = ollama_manager.get_active_processes()
        if processes:
            print("\n活跃的Ollama进程:")
            for pid, info in processes.items():
                running_time = info["running_time"]
                print(f"  进程 {pid}: 运行了 {running_time:.1f} 秒")
                print(f"    命令: {info['command']}")
        else:
            print("\n没有活跃的Ollama进程")
            
        return
    
    # 验证仓库路径
    repo_path = os.path.abspath(args.repo)
    if not os.path.exists(repo_path):
        logger.error(f"仓库路径不存在: {repo_path}")
        return
    
    start_time = time.time()
    logger.info(f"开始分析代码仓库: {repo_path}")
    
    try:
        # 初始化代码分析器
        analyzer = CodeAnalyzer(repo_path)
        analysis_results = analyzer.analyze()
        
        # 初始化知识图谱
        graph_db = GraphDBManager()
        graph_db.store_analysis_results(analysis_results)
        
        # 初始化文档生成器
        global generator
        generator = DocGenerator(analyzer, graph_db)
        
        # 根据命令行参数执行相应操作
        if args.file and args.function:
            # 处理文件路径
            if os.path.isabs(args.file):
                # 如果是绝对路径，确保它在目标仓库内
                if not args.file.startswith(repo_path):
                    logger.error(f"文件路径 {args.file} 不在仓库 {repo_path} 内")
                    return
                # 转换为相对于仓库的路径
                file_path = os.path.relpath(args.file, repo_path)
            else:
                # 如果是相对路径，直接使用
                file_path = args.file
                # 确认文件存在
                full_path = os.path.join(repo_path, file_path)
                if not os.path.exists(full_path):
                    logger.error(f"文件不存在: {full_path}")
                    return
            
            # 生成特定函数的文档
            if args.subprocess:
                result = _generate_with_subprocess(file_path, args.function)
                print(f"函数文档已生成: {args.function}")
            elif args.observe:
                doc = generator.generate_function_doc_observable(file_path, args.function)
                print(f"函数文档已生成: {args.function}")
            else:
                doc = generator.generate_function_doc(file_path, args.function)
                print(f"函数文档已生成: {args.function}")
            
        elif args.file:
            # 处理文件路径
            if os.path.isabs(args.file):
                # 如果是绝对路径，确保它在目标仓库内
                if not args.file.startswith(repo_path):
                    logger.error(f"文件路径 {args.file} 不在仓库 {repo_path} 内")
                    return
                # 转换为相对于仓库的路径
                file_path = os.path.relpath(args.file, repo_path)
            else:
                # 如果是相对路径，直接使用
                file_path = args.file
                # 确认文件存在
                full_path = os.path.join(repo_path, file_path)
                if not os.path.exists(full_path):
                    logger.error(f"文件不存在: {full_path}")
                    return
            
            # 生成特定文件的文档
            if args.subprocess:
                result = _generate_with_subprocess(file_path)
                print(f"文件文档已生成: {file_path}")
            elif args.observe:
                doc = generator.generate_file_doc_observable(file_path)
                print(f"文件文档已生成: {file_path}")
            else:
                doc = generator.generate_file_doc(file_path)
                print(f"文件文档已生成: {file_path}")
            
        elif args.overview:
            # 生成项目概述
            if args.subprocess:
                result = _generate_overview_with_subprocess()
                print("项目概述文档已生成")
            elif args.observe:
                doc = generator.generate_project_overview_observable()
                print("项目概述文档已生成")
            else:
                doc = generator.generate_project_overview()
                print("项目概述文档已生成")
            
        elif args.all:
            # 为所有文件生成文档
            logger.info("开始为所有Python文件生成文档...")
            count = 0
            for file_path in analyzer.files.keys():
                if file_path.endswith('.py'):
                    logger.info(f"处理文件 {count+1}/{len([f for f in analyzer.files if f.endswith('.py')])}: {file_path}")
                    try:
                        if args.subprocess:
                            result = _generate_with_subprocess(file_path)
                        elif args.observe:
                            generator.generate_file_doc_observable(file_path)
                        else:
                            generator.generate_file_doc(file_path)
                        count += 1
                    except Exception as e:
                        logger.error(f"处理文件 {file_path} 时出错: {str(e)}")
            print(f"已为 {count} 个Python文件生成文档")
            
        else:
            # 默认生成项目概述
            if args.subprocess:
                result = _generate_overview_with_subprocess()
                print("项目概述文档已生成")
            elif args.observe:
                doc = generator.generate_project_overview_observable()
                print("项目概述文档已生成")
            else:
                doc = generator.generate_project_overview()
                print("项目概述文档已生成")
        
        elapsed_time = time.time() - start_time
        logger.info(f"文档生成完成，总耗时: {elapsed_time:.2f}秒")
        
    except Exception as e:
        logger.error(f"执行过程中出错: {str(e)}", exc_info=True)
    finally:
        # 清理资源
        if 'graph_db' in locals():
            graph_db.close()

def _generate_with_subprocess(file_path, function_name=None):
    """通过子进程方式生成文档"""
    from config.settings import OUTPUT_DIR, OLLAMA_MODEL
    
    # 准备输出路径
    if function_name:
        output_file = f"{file_path.replace('/', '_')}__{function_name}.md"
    else:
        output_file = f"{file_path.replace('/', '_')}"
    
    output_path = os.path.join(OUTPUT_DIR, output_file)
    
    # 获取文件上下文
    if function_name:
        prompt = generator._prepare_function_doc_prompt(
            file_path, 
            function_name,
            generator.graph_db.get_function_context(file_path, function_name),
            generator.analyzer.get_file_content(file_path)
        )
    else:
        prompt = generator._prepare_file_doc_prompt(
            file_path,
            generator.analyzer.get_context_for_file(file_path)
        )
    
    # 使用进程管理器执行
    print(f"\n==== 使用子进程模式生成文档 ({OLLAMA_MODEL}) ====")
    print(f"生成中，请稍候...")
    
    result = ollama_manager.run_with_timeout(
        model=OLLAMA_MODEL,
        prompt=prompt,
        timeout=300  # 5分钟超时
    )
    
    if result["success"]:
        # 保存文档
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(result["text"])
        print(f"\n文档已保存到: {output_path}")
        return result["text"]
    else:
        print(f"\n生成失败: {result['error']}")
        return None

def _generate_overview_with_subprocess():
    """通过子进程方式生成项目概述"""
    from config.settings import OUTPUT_DIR, OLLAMA_MODEL
    
    # 准备输出路径
    output_path = os.path.join(OUTPUT_DIR, "project_overview.md")
    
    # 统计项目概览
    total_files = len(generator.analyzer.files)
    file_types = {}
    for file in generator.analyzer.files:
        ext = os.path.splitext(file)[1]
        if ext in file_types:
            file_types[ext] += 1
        else:
            file_types[ext] = 1
    
    # 准备提示词
    prompt = f"""
    基于以下项目信息生成一份全面的项目概述文档：
    
    项目统计信息：
    - 总文件数: {total_files}
    - 文件类型分布: {file_types}
    
    重要文件列表:
    {generator._get_important_files()}
    
    主要模块和功能:
    {generator._get_key_modules()}
    
    请生成一份Markdown格式的项目概述文档，包括：
    1. 项目架构概述
    2. 核心模块说明
    3. 关键功能与流程
    4. 技术栈分析
    5. 组件间依赖关系
    
    文档应该简洁明了，重点突出项目的整体结构和主要功能。
    """
    
    # 使用进程管理器执行
    print(f"\n==== 使用子进程模式生成项目概述 ({OLLAMA_MODEL}) ====")
    print(f"生成中，请稍候...")
    
    result = ollama_manager.run_with_timeout(
        model=OLLAMA_MODEL,
        prompt=prompt,
        timeout=300  # 5分钟超时
    )
    
    if result["success"]:
        # 保存文档
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(result["text"])
        print(f"\n项目概述已保存到: {output_path}")
        return result["text"]
    else:
        print(f"\n生成失败: {result['error']}")
        return None

if __name__ == "__main__":
    main()
